# -*- coding: utf-8 -*-
"""
HaaS506 RTU MQTT模块
简化版本 - MQTT通信和消息处理
"""

import ujson as json
import utime
import mqtt
import rtu_config

# MQTT客户端
_mqtt_client = None
_mqtt_connected = False
_last_ping = 0
_control_callback = None
_last_loop_check = 0
_connection_lost_time = 0
_last_connected_time = 0
_reconnect_attempts = 0
_last_mqtt_activity = 0
_mqtt_timeout_threshold = 300  # 5分钟无活动则认为断线
_max_reconnect_attempts = 10  # 增加到10次
_reconnect_delay = 15  # 重连间隔改为15秒
_last_force_reconnect_time = 0  # 上次强制重连时间

def init():
    """初始化MQTT模块"""
    print("[MQTT] MQTT module initialized")
    return True

def connect():
    """连接MQTT服务器"""
    global _mqtt_client, _mqtt_connected, _last_mqtt_activity, _reconnect_attempts
    global _last_connected_time, _connection_lost_time
    
    try:
        if not rtu_config.g_network_ready:
            print("[MQTT] Network not ready")
            return False
        
        # 在弱信号环境下，等待一段时间让网络稳定
        import rtu_network
        if rtu_network.is_signal_weak():
            print("[MQTT] Weak signal detected, waiting for network stability...")
            utime.sleep(2)
        
        mqtt_config = rtu_config.get_mqtt_config()
        
        print("[MQTT] Connecting to {}:{}".format(mqtt_config['broker'], mqtt_config['port']))
        print("[MQTT] Client ID: {}".format(mqtt_config['client_id']))
        
        # 打印更多调试信息
        device_sn = rtu_config.get_device_sn()
        print("[MQTT] Device SN: {}".format(device_sn))
        print("[MQTT] Username: {}".format(mqtt_config.get('username', 'None')))
        
        # 创建MQTT客户端
        _mqtt_client = mqtt.client(mqtt_config['client_id'])
        
        # 设置用户名和密码
        if mqtt_config.get('username') and mqtt_config.get('password'):
            _mqtt_client.username_pw_set(mqtt_config['username'], mqtt_config['password'])
        
        # 连接服务器
        broker = mqtt_config['broker']
        port = int(mqtt_config.get('port', 1883))
        # 在弱信号环境下，增加keepalive时间
        if rtu_network.is_signal_weak():
            keepalive = 120  # 2分钟
            print("[MQTT] Using extended keepalive for weak signal: {}s".format(keepalive))
        else:
            keepalive = int(mqtt_config.get('keepalive', 60))
        
        print("[MQTT] Connecting with broker={}, port={}, keepalive={}".format(broker, port, keepalive))
        
        # 尝试连接，增加异常处理
        try:
            result = _mqtt_client.connect(broker, port, keepalive)
        except OSError as e:
            # 网络错误
            print("[MQTT] Network error during connect: {}".format(e))
            if e.args[0] == 105:  # ENOBUFS
                print("[MQTT] Network buffer exhausted, need to wait for recovery")
            _mqtt_client = None
            return False
        except Exception as e:
            # 其他错误
            print("[MQTT] Connect exception: {}".format(e))
            _mqtt_client = None
            return False
        
        if result != 0:
            print("[MQTT] Connect failed, code: {}".format(result))
            # 根据错误码提供更详细的信息
            if result == -1:
                print("[MQTT] Connection refused - check broker address and port")
            elif result == -2:
                print("[MQTT] Network unreachable")
            elif result == -3:
                print("[MQTT] Connection timeout")
            _mqtt_client = None
            return False
        
        _mqtt_connected = True
        rtu_config.g_mqtt_connected = True
        _last_connected_time = utime.time()
        _last_mqtt_activity = utime.time()
        _reconnect_attempts = 0
        
        # 重置连接丢失时间
        if _connection_lost_time > 0:
            disconnect_duration = utime.time() - _connection_lost_time
            print("[MQTT] Reconnected after {:.1f} seconds".format(disconnect_duration))
            _connection_lost_time = 0
        
        print("[MQTT] Connected successfully")
        
        # 尝试订阅
        try:
            # 设置消息回调
            _mqtt_client.on_subcribe(_on_message_simple)
            
            # 订阅所有主题
            topics = mqtt_config['topics']
            qos = mqtt_config.get('qos', 1)
            
            for topic_name, topic_path in topics.items():
                if topic_name in ['control', 'status']:  # 只订阅控制和状态主题
                    res = _mqtt_client.subscribe(topic_path, qos)
                    print("[MQTT] Subscribe '{}' result: {}".format(topic_name, res))
                    print("[MQTT] Subscribed to {}: {}".format(topic_name, topic_path))
            
        except Exception as e:
            print("[MQTT] Subscribe error: {}".format(e))
            # 即使订阅失败，连接仍然成功
        
        return True
        
    except Exception as e:
        print("[MQTT] Connection failed: {}".format(e))
        _mqtt_connected = False
        rtu_config.g_mqtt_connected = False
        if _connection_lost_time == 0:
            _connection_lost_time = utime.time()
        return False

def disconnect():
    """断开MQTT连接"""
    global _mqtt_client, _mqtt_connected, _connection_lost_time
    
    try:
        if _mqtt_client:
            _mqtt_client.disconnect()
            _mqtt_client = None
        
        _mqtt_connected = False
        rtu_config.g_mqtt_connected = False
        
        # 记录断开时间
        if _connection_lost_time == 0:
            _connection_lost_time = utime.time()
        
        print("[MQTT] Disconnected")
        
    except Exception as e:
        print("[MQTT] Disconnect error: {}".format(e))

def publish_data(data):
    """发布数据"""
    global _mqtt_connected, _connection_lost_time, _mqtt_client
    
    if not _mqtt_connected or not _mqtt_client:
        print("[MQTT] Not connected")
        return False
    
    try:
        mqtt_config = rtu_config.get_mqtt_config()
        topic = mqtt_config['topics']['data']
        payload = json.dumps(data)
        
        # 在弱信号环境下的优化处理
        import rtu_network
        is_weak_signal = rtu_network.is_signal_weak()
        
        # 降低QoS级别以提高成功率
        if is_weak_signal:
            qos = 0  # 使用QoS 0 (最多一次交付)
            print("[MQTT] Using QoS 0 for weak signal condition")
        else:
            qos = mqtt_config.get('qos', 1)
        
        # 在弱信号环境下，先进行快速连接检查
        if is_weak_signal:
            # 快速ping检查
            try:
                # 尝试一个快速的loop操作来验证连接
                check_result = _mqtt_client.loop(10)  # 10ms超时
                if check_result < -1:
                    print("[MQTT] Connection check failed in weak signal")
                    _mqtt_connected = False
                    rtu_config.g_mqtt_connected = False
                    if _connection_lost_time == 0:
                        _connection_lost_time = utime.time()
                    return False
            except:
                pass
        
        result = _mqtt_client.publish(topic, payload, qos)
        
        if result is not None and result >= 0:
            print("[MQTT] Data published to {}".format(topic))
            # 更新活动时间
            global _last_mqtt_activity
            _last_mqtt_activity = utime.time()
            return True
        else:
            print("[MQTT] Publish failed, result: {}".format(result))
            # 发布失败，可能连接已断开
            if result == -1:
                print("[MQTT] Connection may be lost, marking as disconnected")
                _mqtt_connected = False
                rtu_config.g_mqtt_connected = False
                if _connection_lost_time == 0:
                    _connection_lost_time = utime.time()
                
                # 在弱信号环境下，立即清理客户端对象
                if rtu_network.is_signal_weak():
                    print("[MQTT] Weak signal detected, clearing client object")
                    try:
                        _mqtt_client.disconnect()
                    except:
                        pass
                    _mqtt_client = None
            return False
            
    except Exception as e:
        print("[MQTT] Publish error: {}".format(e))
        # 异常情况下也标记为断开
        _mqtt_connected = False
        rtu_config.g_mqtt_connected = False
        if _connection_lost_time == 0:
            _connection_lost_time = utime.time()
        return False

def publish_status(status):
    """发布状态"""
    if not _mqtt_connected or not _mqtt_client:
        print("[MQTT] Not connected")
        return False
    
    try:
        mqtt_config = rtu_config.get_mqtt_config()
        topic = mqtt_config['topics']['status']
        payload = json.dumps(status)
        qos = mqtt_config.get('qos', 1)
        
        result = _mqtt_client.publish(topic, payload, qos)
        
        if result is not None and result >= 0:
            print("[MQTT] Status published to {}".format(topic))
            return True
        else:
            print("[MQTT] Publish status failed")
            return False
            
    except Exception as e:
        print("[MQTT] Publish status error: {}".format(e))
        return False

def publish_ota_status(data):
    """发布OTA状态"""
    return publish_status(data)

def set_control_callback(callback):
    """设置控制消息回调"""
    global _control_callback
    _control_callback = callback

def _on_message_simple(topic, payload):
    """简化的消息回调处理"""
    try:
        print("[MQTT] Message received on topic: {}".format(topic))
        
        # 检查消息大小，防止内存耗尽
        MAX_JSON_SIZE = 4096  # 4KB 限制
        if len(payload) > MAX_JSON_SIZE:
            print("[MQTT] Message too large: {} bytes, max: {}".format(len(payload), MAX_JSON_SIZE))
            return
        
        print("[MQTT] Raw payload: {}".format(payload))
        
        # 解析消息
        try:
            message = json.loads(payload)
            print("[MQTT] Parsed message: {}".format(message))
        except ValueError as e:
            print("[MQTT] Invalid JSON: {}".format(e))
            return
        except MemoryError:
            print("[MQTT] Out of memory parsing JSON")
            return
        except Exception as e:
            print("[MQTT] JSON parse error: {}".format(e))
            return
        
        # 获取MQTT配置以判断主题类型
        mqtt_config = rtu_config.get_mqtt_config()
        
        # 判断主题类型
        if topic == mqtt_config['topics']['control']:
            print("[MQTT] Control/Write command received")
            
            # 如果设置了控制回调，调用它
            if _control_callback:
                try:
                    _control_callback(message)
                    
                    # 发送状态反馈
                    status_response = {
                        "cmd": message.get("command", message.get("action", "unknown")),
                        "status": "success",
                        "timestamp": utime.time(),
                        "message": "Command executed successfully"
                    }
                    publish_status(status_response)
                    
                except Exception as e:
                    print("[MQTT] Control callback error: {}".format(e))
                    
                    # 发送错误状态反馈
                    status_response = {
                        "cmd": message.get("command", message.get("action", "unknown")),
                        "status": "error",
                        "timestamp": utime.time(),
                        "message": "Command failed: {}".format(str(e))
                    }
                    publish_status(status_response)
            else:
                print("[MQTT] No control callback set")
                
        elif topic == mqtt_config['topics']['status']:
            print("[MQTT] Status message received: {}".format(message))
            # 状态消息通常是服务器对我们发送状态的确认，可以不处理
            
        else:
            print("[MQTT] Unknown topic: {}".format(topic))
                
    except Exception as e:
        print("[MQTT] Message processing error: {}".format(e))

def _on_message(topic, msg):
    """消息回调处理"""
    try:
        # 检查消息大小
        MAX_JSON_SIZE = 4096
        if len(msg) > MAX_JSON_SIZE:
            print("[MQTT] Message too large: {} bytes".format(len(msg)))
            return
        
        topic_str = topic.decode('utf-8')
        msg_str = msg.decode('utf-8')
        
        print("[MQTT] Received message from {}: {}".format(topic_str, msg_str))
        
        # 解析消息
        try:
            message = json.loads(msg_str)
        except ValueError as e:
            print("[MQTT] Invalid JSON: {}".format(e))
            return
        except MemoryError:
            print("[MQTT] Out of memory parsing JSON")
            return
        
        # 如果设置了控制回调，调用它
        if _control_callback:
            _control_callback(message)
            
    except Exception as e:
        print("[MQTT] Message processing error: {}".format(e))

def check_connection():
    """检查连接状态"""
    global _mqtt_connected
    
    if not _mqtt_client:
        _mqtt_connected = False
        rtu_config.g_mqtt_connected = False
        return False
    
    # 这里可以添加ping检查
    return _mqtt_connected

def is_connected():
    """是否已连接"""
    return _mqtt_connected

def get_mqtt_status():
    """获取MQTT状态"""
    return {
        "connected": _mqtt_connected,
        "last_ping": _last_ping
    }

def process_loop():
    """处理MQTT消息循环 - 被主循环定期调用"""
    global _mqtt_client, _mqtt_connected, _last_loop_check, _last_mqtt_activity
    global _connection_lost_time, _last_ping
    
    if not _mqtt_client:
        return
    
    # 如果标记为未连接，直接返回
    if not _mqtt_connected:
        return
    
    try:
        current_time = utime.time()
        
        # 每5秒调用一次loop，检查是否有消息
        if current_time - _last_loop_check >= 5:
            _last_loop_check = current_time
            
            # 调用loop处理消息，使用较短的超时时间
            try:
                ret = _mqtt_client.loop(100)  # 100ms超时
                
                # MicroPython MQTT库的loop()可能返回-1表示没有消息，而不是断开连接
                # 只有在真正的连接错误时才断开
                if ret < -1:  # 严重错误
                    print("[MQTT] Loop check failed, severe error (ret={})".format(ret))
                    _mqtt_connected = False
                    rtu_config.g_mqtt_connected = False
                    if _connection_lost_time == 0:
                        _connection_lost_time = utime.time()
                else:
                    # 即使ret=-1，也认为连接正常（可能只是没有新消息）
                    # 记录MQTT活动时间
                    _last_mqtt_activity = current_time
                    
                    # 如果之前标记为断开，但现在能正常loop，说明连接恢复了
                    if not _mqtt_connected:
                        _mqtt_connected = True
                        rtu_config.g_mqtt_connected = True
                        print("[MQTT] Connection restored during loop check")
                        
            except Exception as e:
                # 真正的异常才认为断开
                print("[MQTT] Loop exception: {}".format(e))
                _mqtt_connected = False
                rtu_config.g_mqtt_connected = False
                if _connection_lost_time == 0:
                    _connection_lost_time = utime.time()
        
        # 注释掉ping功能，因为MicroPython的MQTT可能不支持
        # MQTT keepalive会自动处理心跳
        # if current_time - _last_ping >= 30:
        #     _last_ping = current_time
        #     # MicroPython mqtt库可能不支持ping方法
            
    except Exception as e:
        print("[MQTT] Process loop error: {}".format(e))
        # 发生异常，标记为断开
        _mqtt_connected = False
        rtu_config.g_mqtt_connected = False
        if _connection_lost_time == 0:
            _connection_lost_time = utime.time()

def quick_check():
    """快速检查消息 - 用于需要及时响应的场景"""
    global _mqtt_client, _mqtt_connected, _connection_lost_time
    
    if not _mqtt_client or not _mqtt_connected:
        return
    
    try:
        # 使用非常短的超时快速检查
        ret = _mqtt_client.loop(10)  # 10ms超时
        
        # 只有严重错误才认为断开
        if ret < -1:
            _mqtt_connected = False
            rtu_config.g_mqtt_connected = False
            if _connection_lost_time == 0:
                _connection_lost_time = utime.time()
            
    except Exception as e:
        pass  # 快速检查失败不打印错误

def is_mqtt_timeout():
    """检查MQTT是否超时"""
    if not _mqtt_connected:
        return False
    
    if _last_mqtt_activity == 0:
        return False
    
    inactive_time = utime.time() - _last_mqtt_activity
    return inactive_time > _mqtt_timeout_threshold

def get_mqtt_stats():
    """获取MQTT连接统计信息"""
    current_time = utime.time()
    
    if _mqtt_connected:
        uptime = current_time - _last_connected_time if _last_connected_time > 0 else 0
        downtime = 0
        inactive_time = current_time - _last_mqtt_activity if _last_mqtt_activity > 0 else 0
    else:
        uptime = 0
        downtime = current_time - _connection_lost_time if _connection_lost_time > 0 else 0
        inactive_time = 0
    
    return {
        "connected": _mqtt_connected,
        "uptime": uptime,
        "downtime": downtime,
        "inactive_time": inactive_time,
        "reconnect_attempts": _reconnect_attempts,
        "timeout_threshold": _mqtt_timeout_threshold
    }

def attempt_reconnect():
    """尝试重新连接MQTT"""
    global _reconnect_attempts, _mqtt_client
    
    # 在弱信号环境下，适当调整重试次数
    import rtu_network
    if rtu_network.is_signal_weak():
        max_attempts = 5  # 弱信号时改为5次
        print("[MQTT] Weak signal detected, limiting reconnect attempts to {}".format(max_attempts))
    else:
        max_attempts = _max_reconnect_attempts
    
    if _reconnect_attempts >= max_attempts:
        print("[MQTT] Max reconnect attempts reached, will wait for network recovery")
        # 清理客户端对象，下次重新创建
        if _mqtt_client:
            try:
                _mqtt_client.disconnect()
            except:
                pass
            _mqtt_client = None
        return False
    
    print("[MQTT] Attempting reconnect {}/{}".format(_reconnect_attempts + 1, _max_reconnect_attempts))
    
    # 清理现有连接
    if _mqtt_client:
        try:
            _mqtt_client.disconnect()
        except:
            pass
        _mqtt_client = None
    
    # 使用指数退避算法，在弱信号环境下减少重连频率
    base_delay = 3  # 基础延迟改为3秒
    max_delay = 120  # 最大延迟改为2分钟
    
    # 指数退避：3, 6, 12, 24, 48, 96, 120秒
    delay = min(base_delay * (2 ** _reconnect_attempts), max_delay)
    
    # 添加随机抖动，避免多设备同时重连
    import urandom
    jitter = urandom.getrandbits(3)  # 0-7秒随机
    delay = delay + jitter
    
    print("[MQTT] Waiting {} seconds before reconnect (attempt {})...".format(delay, _reconnect_attempts + 1))
    utime.sleep(delay)
    
    # 尝试重新连接
    if connect():
        print("[MQTT] Reconnect successful")
        _reconnect_attempts = 0  # 成功后重置计数
        return True
    else:
        _reconnect_attempts += 1
        print("[MQTT] Reconnect failed, attempt {}/{}".format(_reconnect_attempts, _max_reconnect_attempts))
        return False

def force_reconnect():
    """强制重新连接MQTT"""
    global _reconnect_attempts, _mqtt_client, _last_force_reconnect_time
    
    # 防止过于频繁的强制重连
    current_time = utime.time()
    if _last_force_reconnect_time > 0:
        if current_time - _last_force_reconnect_time < 30:  # 改为30秒
            print("[MQTT] Force reconnect called too frequently, skipping")
            return False
    _last_force_reconnect_time = current_time
    
    print("[MQTT] Force reconnecting...")
    
    # 先检查网络状态
    if not rtu_config.g_network_ready:
        print("[MQTT] Network not ready, cannot reconnect")
        return False
    
    # 清理现有客户端
    if _mqtt_client:
        try:
            _mqtt_client.disconnect()
        except:
            pass
        _mqtt_client = None
    
    # 不重置计数器，让它继续累加
    # _reconnect_attempts = 0  # 注释掉这行
    
    # 如果已经达到最大重试次数，重置计数器
    if _reconnect_attempts >= _max_reconnect_attempts:
        print("[MQTT] Resetting reconnect counter for force reconnect")
        _reconnect_attempts = 0
    
    # 等待一下让网络稳定
    utime.sleep(2)
    
    # 强制重连
    return attempt_reconnect()

def on_network_recovery():
    """网络恢复时的回调函数"""
    global _reconnect_attempts, _last_force_reconnect_time
    
    print("[MQTT] Network recovered, attempting MQTT reconnection...")
    
    # 重置重连计数和强制重连时间
    _reconnect_attempts = 0
    _last_force_reconnect_time = 0  # 允许立即强制重连
    
    # 延迟一点时间再重连，确保网络稳定
    utime.sleep(3)
    
    # 尝试重新连接
    if connect():
        print("[MQTT] MQTT reconnected after network recovery")
        return True
    else:
        print("[MQTT] MQTT reconnection failed after network recovery")
        return False

def is_long_disconnected(threshold_minutes=5):
    """检查MQTT是否长时间断开"""
    if _mqtt_connected:
        return False
    
    if _connection_lost_time == 0:
        return False
    
    disconnected_time = utime.time() - _connection_lost_time
    return disconnected_time > (threshold_minutes * 60)

def test_mqtt():
    """测试MQTT功能"""
    print("[MQTT] Testing MQTT functionality...")
    
    # 测试连接
    if connect():
        print("[MQTT] Connection test passed")
        
        # 测试发布
        test_data = {
            "test": "data",
            "timestamp": utime.time()
        }
        
        if publish_data(test_data):
            print("[MQTT] Publish test passed")
        else:
            print("[MQTT] Publish test failed")
        
        # 断开连接
        disconnect()
    else:
        print("[MQTT] Connection test failed")
    
    print("[MQTT] MQTT test completed")